import os
import sys
import subprocess
import argparse
from pathlib import Path
import shutil
import platform

def setup_vs_environment():
    """配置 Visual Studio 开发环境变量"""
    print("正在配置 Visual Studio 环境...")
    
    # 尝试找到 Visual Studio 的安装路径
    possible_paths = [
        r"C:\Program Files (x86)\Microsoft Visual Studio\2019\Community\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files\Microsoft Visual Studio\2019\Community\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files (x86)\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files (x86)\Microsoft Visual Studio\2017\Community\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files\Microsoft Visual Studio\2017\Community\VC\Auxiliary\Build\vcvarsall.bat",
        # 专业版和企业版的路径
        r"C:\Program Files (x86)\Microsoft Visual Studio\2019\Professional\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files\Microsoft Visual Studio\2019\Professional\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files (x86)\Microsoft Visual Studio\2019\Enterprise\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files\Microsoft Visual Studio\2019\Enterprise\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files (x86)\Microsoft Visual Studio\2022\Professional\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files\Microsoft Visual Studio\2022\Professional\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files (x86)\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvarsall.bat",
    ]
    
    vcvarsall_path = None
    for path in possible_paths:
        if os.path.exists(path):
            vcvarsall_path = path
            break
    
    if not vcvarsall_path:
        print("警告: 无法找到 Visual Studio 的 vcvarsall.bat 文件")
        return None
    
    print(f"找到 vcvarsall.bat: {vcvarsall_path}")
    
    # 确定要使用的架构
    arch = "x64" if platform.machine().endswith('64') else "x86"
    
    # 运行 vcvarsall.bat 并获取环境变量
    cmd = f'"{vcvarsall_path}" {arch} && set'
    print(f"运行命令获取环境变量: {cmd}")
    
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    output, error = process.communicate()
    
    if process.returncode != 0:
        print(f"警告: vcvarsall.bat 执行失败: {error.decode('utf-8', errors='replace')}")
        return None
    
    # 解析环境变量
    env = {}
    for line in output.decode('utf-8', errors='replace').splitlines():
        if '=' in line:
            key, value = line.split('=', 1)
            env[key] = value
    
    # 输出关键环境变量以便调试
    important_vars = ['PATH', 'INCLUDE', 'LIB', 'LIBPATH', 'WindowsSdkDir', 'WindowsSDKVersion']
    for var in important_vars:
        if var in env:
            print(f"设置环境变量 {var}: {env[var][:50]}..." if len(env[var]) > 50 else f"设置环境变量 {var}: {env[var]}")
    
    # 确保PATH环境变量被正确设置
    if 'PATH' not in env or not env['PATH']:
        print("错误: PATH环境变量未在vcvarsall.bat中设置")
        # 尝试使用环境变量获取PATH
        if 'Path' in env:  # Windows环境变量名称可能是 'Path' 而非 'PATH'
            print("使用 'Path' 环境变量替代")
            env['PATH'] = env['Path']
        elif os.environ.get('PATH'):
            print("从当前进程环境中获取PATH")
            env['PATH'] = os.environ.get('PATH')
        
        # 尝试查找cl.exe的位置
        vs_dir = os.path.dirname(os.path.dirname(os.path.dirname(vcvarsall_path)))
        cl_paths = []
        
        # 搜索常见位置
        for root, dirs, files in os.walk(vs_dir):
            if 'cl.exe' in files:
                cl_paths.append(os.path.join(root, 'cl.exe'))
        
        if cl_paths:
            print(f"找到以下cl.exe: {cl_paths}")
            if 'PATH' not in env:
                env['PATH'] = ''
            
            # 添加cl.exe所在目录到PATH
            for cl_path in cl_paths:
                cl_dir = os.path.dirname(cl_path)
                if cl_dir not in env['PATH']:
                    env['PATH'] = f"{cl_dir};{env['PATH']}"
                    print(f"添加 {cl_dir} 到PATH")
        else:
            print("错误: 无法找到cl.exe")
    else:
        # 检查PATH中是否包含cl.exe
        cl_found = False
        path_entries = env['PATH'].split(';')
        for entry in path_entries:
            cl_path = os.path.join(entry, 'cl.exe')
            if os.path.exists(cl_path):
                print(f"PATH中已包含cl.exe: {cl_path}")
                cl_found = True
                break
        
        if not cl_found:
            print("警告: PATH中未找到cl.exe, 尝试手动查找...")
            
            # 查找Visual Studio的安装路径
            vs_dir = os.path.dirname(os.path.dirname(os.path.dirname(vcvarsall_path)))
            cl_paths = []
            
            # 搜索常见位置
            for root, dirs, files in os.walk(vs_dir):
                if 'cl.exe' in files:
                    cl_paths.append(os.path.join(root, 'cl.exe'))
            
            if cl_paths:
                print(f"找到以下cl.exe: {cl_paths}")
                # 添加cl.exe所在目录到PATH
                for cl_path in cl_paths:
                    cl_dir = os.path.dirname(cl_path)
                    if cl_dir not in env['PATH']:
                        env['PATH'] = f"{cl_dir};{env['PATH']}"
                        print(f"添加 {cl_dir} 到PATH")
            else:
                print("错误: 无法找到cl.exe")
            
    return env

def run_command(cmd, cwd=None, env=None):
    """运行命令并实时输出结果"""
    print(f"执行命令: {' '.join(cmd)}")
    
    # 合并当前环境变量和提供的环境变量
    merged_env = os.environ.copy()
    if env:
        print(f"合并环境变量...")
        # 确保我们覆盖现有的环境变量，而不是简单地添加
        for key, value in env.items():
            merged_env[key] = value
            
        # 确认关键环境变量确实被传递
        if 'PATH' in env:
            print(f"PATH环境变量已设置，长度: {len(env['PATH'])}")
    
    process = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=False,  # 使用二进制模式
        cwd=cwd,
        env=merged_env
    )
    
    while True:
        output = process.stdout.readline()
        if output == b'' and process.poll() is not None:
            break
        if output:
            try:
                # 尝试用 utf-8 解码，忽略错误
                decoded_output = output.decode('utf-8', errors='replace')
                print(decoded_output.strip())
            except Exception as e:
                # 解码失败时显示原始字节
                print(f"[无法解码的输出]: {output}")
    
    return process.poll()

def check_generator_exists(generator):
    """检查指定的生成器是否存在"""
    cmd = ["cmake", "--help"]
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=False)
    output, _ = process.communicate()
    output_str = output.decode('utf-8', errors='replace')
    
    # 检查生成器是否在帮助输出中列出
    return generator in output_str

def check_ninja_works():
    """检查Ninja是否可用"""
    try:
        # 尝试运行ninja --version
        process = subprocess.Popen(["ninja", "--version"], 
                                  stdout=subprocess.PIPE, 
                                  stderr=subprocess.PIPE)
        _, _ = process.communicate()
        return process.returncode == 0
    except Exception:
        return False

def build_cmake_project(project_dir, build_type="Release", build_dir=None, clean=False, generator=None):
    """使用 CMake 构建项目并生成 compile_commands.json"""
    project_path = Path(project_dir)
    
    # 使用自定义构建目录或默认的 "build" 目录
    if build_dir:
        build_dir = Path(build_dir)
    else:
        build_dir = project_path / "build"
    
    # 如果指定了 clean 参数，清理构建目录
    if clean and build_dir.exists():
        print(f"清理构建目录: {build_dir}")
        shutil.rmtree(build_dir)
    
    # 创建构建目录（如果不存在）
    os.makedirs(build_dir, exist_ok=True)
    
    # 设置 Visual Studio 环境
    vs_env = setup_vs_environment()
    if not vs_env:
        print("警告: 未配置 Visual Studio 环境，可能会影响 CMake 查找编译器的能力")
    else:
        print("成功配置 Visual Studio 环境变量")
    
    # 保存vcvarsall_path用于后续判断
    vcvarsall_path = None
    for path in [
        r"C:\Program Files (x86)\Microsoft Visual Studio\2019\Community\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files\Microsoft Visual Studio\2019\Community\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files (x86)\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files (x86)\Microsoft Visual Studio\2017\Community\VC\Auxiliary\Build\vcvarsall.bat",
        r"C:\Program Files\Microsoft Visual Studio\2017\Community\VC\Auxiliary\Build\vcvarsall.bat",
    ]:
        if os.path.exists(path):
            vcvarsall_path = path
            break
    
    # 确定当前系统架构
    is_64bit = platform.machine().endswith('64')
    target_arch = "x64" if is_64bit else "x86"
    print(f"检测到系统架构: {target_arch}")
    
    # 查找资源编译器(rc.exe)
    rc_path = None
    if vs_env:
        # 检查Windows SDK路径
        sdk_dir = None
        for key in ['WindowsSdkDir', 'WindowsSDKDir', 'WindowsSDK_ExecutablePath_x86', 'WindowsSDK_ExecutablePath_x64']:
            if key in vs_env:
                sdk_base = vs_env[key]
                print(f"找到Windows SDK路径: {sdk_base}")
                
                # 常见rc.exe位置
                possible_rc_paths = [
                    os.path.join(sdk_base, 'bin', target_arch, 'rc.exe'),  # 优先选择匹配目标架构的rc.exe
                    os.path.join(sdk_base, 'bin', 'rc.exe')
                ]
                
                for rc in possible_rc_paths:
                    if os.path.exists(rc):
                        rc_path = rc
                        print(f"找到资源编译器: {rc_path}")
                        
                        # 将资源编译器路径添加到PATH
                        rc_dir = os.path.dirname(rc_path)
                        if rc_dir not in vs_env.get('PATH', ''):
                            vs_env['PATH'] = f"{rc_dir};{vs_env.get('PATH', '')}"
                            print(f"添加资源编译器路径到PATH: {rc_dir}")
                        break
                
                if rc_path:
                    break
        
        # 如果在SDK路径中没找到，尝试在常见位置查找
        if not rc_path:
            print("在SDK路径中未找到rc.exe，尝试在常见位置查找...")
            windows_kits_paths = [
                r"C:\Program Files (x86)\Windows Kits\10\bin",
                r"C:\Program Files\Windows Kits\10\bin",
                r"C:\Program Files (x86)\Windows Kits\8.1\bin",
                r"C:\Program Files\Windows Kits\8.1\bin"
            ]
            
            for kit_path in windows_kits_paths:
                if os.path.exists(kit_path):
                    print(f"检查Windows Kit路径: {kit_path}")
                    # 检查不同架构目录，优先选择与目标架构匹配的
                    arch_priority = [target_arch]  # 优先检查目标架构
                    if target_arch == "x64":
                        arch_priority.append("x86")  # 其次检查x86
                    else:
                        arch_priority.append("x64")  # 其次检查x64
                    arch_priority.append("")  # 最后检查无架构目录
                    
                    for arch in arch_priority:
                        check_path = os.path.join(kit_path, arch, 'rc.exe')
                        if os.path.exists(check_path):
                            rc_path = check_path
                            rc_dir = os.path.dirname(rc_path)
                            
                            if rc_dir not in vs_env.get('PATH', ''):
                                vs_env['PATH'] = f"{rc_dir};{vs_env.get('PATH', '')}"
                                print(f"添加资源编译器路径到PATH: {rc_dir}")
                            
                            print(f"找到资源编译器: {rc_path}")
                            break
                
                if rc_path:
                    break
            
            # 如果仍然找不到，尝试搜索整个Visual Studio目录
            if not rc_path:
                print("在Windows Kits中未找到rc.exe，尝试搜索Visual Studio目录...")
                vs_dir = os.path.dirname(os.path.dirname(os.path.dirname(vcvarsall_path)))
                
                for root, dirs, files in os.walk(vs_dir):
                    if 'rc.exe' in files:
                        rc_path = os.path.join(root, 'rc.exe')
                        rc_dir = os.path.dirname(rc_path)
                        
                        if rc_dir not in vs_env.get('PATH', ''):
                            vs_env['PATH'] = f"{rc_dir};{vs_env.get('PATH', '')}"
                            print(f"添加资源编译器路径到PATH: {rc_dir}")
                        
                        print(f"找到资源编译器: {rc_path}")
                        break
    
    # 尝试使用支持 compile_commands.json 的生成器
    # 按优先级顺序检查生成器
    available_generator = None
    for generator in ["Ninja", "NMake Makefiles", "MinGW Makefiles"]:
        if check_generator_exists(generator):
            # 对于Ninja，额外检查是否可用
            if generator == "Ninja" and not check_ninja_works():
                print(f"Ninja生成器存在但无法运行，跳过")
                continue
            available_generator = generator
            print(f"使用 {generator} 生成器")
            break
    
    # 配置 CMake 项目，启用 compile_commands.json 生成
    cmake_config_cmd = [
        "cmake", 
        ".."
    ]
    
    # 添加生成器选项
    if generator:
        # 使用用户指定的生成器
        print(f"使用用户指定的生成器: {generator}")
        cmake_config_cmd.extend(["-G", generator])
    elif available_generator:
        # 使用自动检测到的生成器
        cmake_config_cmd.extend(["-G", available_generator])
    
    # 添加其他配置
    cmake_config_cmd.extend([
        f"-DCMAKE_BUILD_TYPE={build_type}", 
        "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON"
    ])
    
    # 查找cl.exe路径以便直接指定
    cl_path = None
    if vs_env and 'PATH' in vs_env:
        for entry in vs_env['PATH'].split(';'):
            if entry and os.path.exists(entry):
                potential_cl_path = os.path.join(entry, 'cl.exe')
                if os.path.exists(potential_cl_path):
                    # 检查是否与目标架构匹配
                    if target_arch == "x64" and ("x64" in entry.lower() or "amd64" in entry.lower()):
                        cl_path = potential_cl_path
                        print(f"找到匹配x64架构的C++编译器: {cl_path}")
                        break
                    elif target_arch == "x86" and ("x86" in entry.lower() and "x64" not in entry.lower()):
                        cl_path = potential_cl_path
                        print(f"找到匹配x86架构的C++编译器: {cl_path}")
                        break
        
        # 如果没有找到匹配的编译器，使用第一个可用的
        if not cl_path:
            for entry in vs_env['PATH'].split(';'):
                if entry and os.path.exists(entry):
                    potential_cl_path = os.path.join(entry, 'cl.exe')
                    if os.path.exists(potential_cl_path):
                        cl_path = potential_cl_path
                        print(f"未找到匹配架构的编译器，使用: {cl_path}")
                        break
    
    # 如果没有找到了cl.exe，尝试手动查找
    if not cl_path:
        print("在PATH环境变量中未找到编译器，尝试手动查找...")
        vs_paths = [
            r"C:\Program Files (x86)\Microsoft Visual Studio",
            r"C:\Program Files\Microsoft Visual Studio"
        ]
        
        for vs_path in vs_paths:
            if os.path.exists(vs_path):
                # 只搜索特定的目录以加快速度
                for year in ["2022", "2019", "2017"]:
                    for edition in ["Community", "Professional", "Enterprise"]:
                        msvc_path = os.path.join(vs_path, year, edition, "VC", "Tools", "MSVC")
                        if os.path.exists(msvc_path):
                            # 查找最新版本的MSVC工具
                            try:
                                versions = [d for d in os.listdir(msvc_path) if os.path.isdir(os.path.join(msvc_path, d))]
                                if versions:
                                    latest = sorted(versions)[-1]  # 获取最新版本
                                    bin_path = os.path.join(msvc_path, latest, "bin", "Hostx64", "x64")
                                    cl_exe = os.path.join(bin_path, "cl.exe")
                                    
                                    if os.path.exists(cl_exe):
                                        cl_path = cl_exe
                                        print(f"手动查找到编译器: {cl_path}")
                                        
                                        # 如果未设置环境变量，创建一个基本的环境
                                        if not vs_env:
                                            vs_env = os.environ.copy()
                                        
                                        # 将找到的编译器路径添加到PATH
                                        if 'PATH' in vs_env:
                                            vs_env['PATH'] = f"{bin_path};{vs_env['PATH']}"
                                        else:
                                            vs_env['PATH'] = bin_path
                                        
                                        print(f"已将编译器路径添加到环境变量: {bin_path}")
                                        break
                            except Exception as e:
                                print(f"查找MSVC版本时出错: {e}")
            
            if cl_path:
                break
    
    # 如果找到了cl.exe，直接在CMake中指定编译器路径
    if cl_path:
        # 将Windows路径转换为CMake可接受的格式（使用正斜杠）
        cmake_cl_path = cl_path.replace('\\', '/')
        cmake_config_cmd.extend([
            f"-DCMAKE_CXX_COMPILER={cmake_cl_path}",
            f"-DCMAKE_C_COMPILER={cmake_cl_path}"
        ])
        print(f"显式指定编译器路径: {cl_path}")
        
    # 如果找到了rc.exe，显式指定资源编译器路径
    if rc_path:
        cmake_rc_path = rc_path.replace('\\', '/')
        cmake_config_cmd.extend([
            f"-DCMAKE_RC_COMPILER={cmake_rc_path}"
        ])
        print(f"显式指定资源编译器路径: {rc_path}")
    
    # 检查是否存在编译器
    if vs_env:
        # 检查关键编译器路径是否在PATH中
        path_entries = vs_env.get('PATH', '').split(';')
        compiler_found = False
        for entry in path_entries:
            if entry and os.path.exists(entry):
                cl_exe_path = os.path.join(entry, 'cl.exe')
                if os.path.exists(cl_exe_path):
                    print(f"找到 MSVC 编译器: {cl_exe_path}")
                    compiler_found = True
                    break
        
        if not compiler_found:
            print("警告: 在环境变量PATH中未找到MSVC编译器(cl.exe)")
    
    # 尝试使用 VS 环境变量运行 CMake
    result = run_command(cmake_config_cmd, cwd=build_dir, env=vs_env)
    
    # 如果配置失败，尝试其他解决方案
    if result != 0:
        print(f"CMake 配置失败，退出码: {result}")
        
        # 调试信息
        print("\n===== 调试信息 =====")
        print("检查CMake是否能找到编译器...")
        run_command(["cmake", "--version"], env=vs_env)
        
        # 如果使用非默认生成器失败，尝试使用默认生成器
        if generator or available_generator:
            print("\n尝试使用默认生成器...")
            cmake_config_cmd = [
                "cmake", 
                "..",
                f"-DCMAKE_BUILD_TYPE={build_type}", 
                "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON"
            ]
            
            # 如果找到了编译器，也为默认生成器指定编译器
            if cl_path:
                cmake_cl_path = cl_path.replace('\\', '/')
                cmake_config_cmd.extend([
                    f"-DCMAKE_CXX_COMPILER={cmake_cl_path}",
                    f"-DCMAKE_C_COMPILER={cmake_cl_path}"
                ])
            
            result = run_command(cmake_config_cmd, cwd=build_dir, env=vs_env)
            if result != 0:
                print(f"使用默认生成器也失败，退出码: {result}")
                
                # 如果使用默认生成器失败，尝试使用Visual Studio生成器
                print("\n尝试使用Visual Studio生成器...")
                
                # 确定Visual Studio版本
                vs_year = None
                vs_version = None
                
                if vcvarsall_path:
                    if "2022" in vcvarsall_path:
                        vs_year = "2022"
                        vs_version = "17"
                    elif "2019" in vcvarsall_path:
                        vs_year = "2019"
                        vs_version = "16"
                    elif "2017" in vcvarsall_path:
                        vs_year = "2017"
                        vs_version = "15"
                        
                if vs_year and vs_version:
                    generator = f"Visual Studio {vs_version} {vs_year}"
                    vs_arch = "x64" if target_arch == "x64" else "Win32"
                    
                    cmake_config_cmd = [
                        "cmake", 
                        "..",
                        f"-G", generator,
                        f"-A", vs_arch,
                        f"-DCMAKE_EXPORT_COMPILE_COMMANDS=ON"
                    ]
                    
                    print(f"使用Visual Studio生成器: {generator}, 架构: {vs_arch}")
                    result = run_command(cmake_config_cmd, cwd=build_dir, env=vs_env)
                    if result != 0:
                        print(f"使用Visual Studio生成器也失败，退出码: {result}")
                        return False
                else:
                    print(f"无法确定Visual Studio版本，无法使用Visual Studio生成器")
                    return False
            return True  # 如果成功则返回True
        else:
            return False
    
    # 构建项目
    cmake_build_cmd = ["cmake", "--build", ".", f"--config={build_type}"]
    result = run_command(cmake_build_cmd, cwd=build_dir, env=vs_env)
    if result != 0:
        print(f"CMake 构建失败，退出码: {result}")
        return False
    
    # 检查 compile_commands.json 是否生成
    compile_commands = build_dir / "compile_commands.json"
    
    # 如果是 VS 生成器，可能在子目录中生成
    if not compile_commands.exists():
        # 搜索可能的位置
        potential_paths = list(build_dir.glob("**/compile_commands.json"))
        if potential_paths:
            # 复制到根构建目录
            shutil.copy(str(potential_paths[0]), str(compile_commands))
            print(f"已找到并复制 compile_commands.json 到 {compile_commands}")
        else:
            # 使用 CMAKE_EXPORT_COMPILATION_DATABASE 模块生成
            print("手动尝试生成 compile_commands.json...")
            cmd = ["cmake", "-E", "cmake_export_compile_commands", "."]
            run_command(cmd, cwd=build_dir, env=vs_env)
    
    # 再次检查
    if compile_commands.exists():
        print(f"成功生成 compile_commands.json: {compile_commands}")
    else:
        print("警告: compile_commands.json 未能生成")
        print("提示: 可以尝试安装 Ninja 或 NMake 以获得更好的支持")
    
    return True

def main():
    parser = argparse.ArgumentParser(description="自动化 CMake 构建流程")
    parser.add_argument("project_dir", help="CMake 项目目录路径")
    parser.add_argument("--build-type", default="Release", help="构建类型 (Debug/Release)")
    parser.add_argument("--build-dir", help="自定义构建目录（可选）")
    parser.add_argument("--clean", action="store_true", help="在构建前清理build目录")
    parser.add_argument("--generator", help="指定CMake生成器（默认自动选择）")
    
    args = parser.parse_args()
    
    success = build_cmake_project(
        args.project_dir, 
        args.build_type,
        build_dir=args.build_dir,
        clean=args.clean,
        generator=args.generator
    )
    return 0 if success else 1

if __name__ == "__main__":
    sys.exit(main()) 